import numpy as np
import tmm
# from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import time

# Define Constants
N_SIO2 = 1.45
N_SI3N4 = 2
min_t = 10
max_t = 100
unit = 'nm/THz'
num_layers = 20
t_sample_num = 2
freq_start = 450
freq_end = 750
num_freq = 200
n_top = 1
n_bottom = 1

n_sio2 = np.ones(num_layers // 2) * N_SIO2
n_si3n4 = np.ones(num_layers // 2) * N_SI3N4
n_array = np.column_stack((n_sio2, n_si3n4)).flatten()
if num_layers % 2 != 0:
    n_array = np.append(n_array, N_SIO2)



def generate_t_sample_array_normalized():
    """
    generate sample array.
    Returns:
        Iterator, give each sample of the thickness combination, normalized in [0,1]
    """
    t_sample = np.linspace(0, 1, t_sample_num + 2)

    for i in range(t_sample_num ** num_layers):
        binary_repr = np.base_repr(i, t_sample_num).zfill(num_layers)
        t_sample_array_elem = np.array([t_sample[1:-1][int(x)] for x in binary_repr])
        yield t_sample_array_elem

def generate_t_sample_array_normalized_rand():
    """
    generate sample array.
        Iterator, give each sample of the thickness combination, normalized in [0,1]
    """

    return np.random.rand(t_sample_num ** num_layers, num_layers)

def calculate_transmission_parallel(t_sample_array_elem):
    """
    calculate transmission spectrum for each thickness combination.
    Returns:
        tran: array, transmission spectrum.
    """
    try:
        t_sample_array_elem = t_sample_array_elem * (max_t - min_t) + min_t
        # convert normalized thickness to real thickness

        _, tran, _ = np.abs(tmm.TMM(t_sample_array_elem, n_array, unit=unit, n_top=n_top, n_bottom=n_bottom)
                            .calculate_spectrum(freq_start, freq_end, num_freq)) ** 2

        return tran
    except Exception as e:
        print(f"error in the calculation of the calculation of the transmission spectrum: {e}")
        return None


def data_generator():


    t_sample_array = list(generate_t_sample_array_normalized_rand())
    print('t_sample_array generated')

    with multiprocessing.Pool(processes=60) as pool:
        tran_array = list(pool.map(calculate_transmission_parallel,t_sample_array))

    # 用try-except捕获np.save可能的异常
    try:
        time_stamp = time.strftime('%Y-%m-%d_%H-%M', time.localtime())
        np.save(f'train_set/t_array_normalized_{num_layers}layers_{t_sample_num}samp_{time_stamp}_rand.npy',
                np.array(t_sample_array))
        np.save(f'train_set/tran_array_normalized_{num_layers}layers_{t_sample_num}samp_{time_stamp}_rand.npy',
                np.array(tran_array))
    except Exception as e:
        print(f"error in saving the array file: {e}")

    return


if __name__ == '__main__':
    data_generator()